{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp models.tsmixer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TSMixer\n",
    "> Time-Series Mixer (`TSMixer`) is a MLP-based multivariate time-series forecasting model. `TSMixer` jointly learns temporal and cross-sectional representations of the time-series by repeatedly combining time- and feature information using stacked mixing layers. A mixing layer consists of a sequential time- and feature Multi Layer Perceptron (`MLP`). Note: this model cannot handle exogenous inputs. If you want to use additional exogenous inputs, use `TSMixerx`.\n",
    "\n",
    "<br><br>**References**<br>-[Chen, Si-An, Chun-Liang Li, Nate Yoder, Sercan O. Arik, and Tomas Pfister (2023). \"TSMixer: An All-MLP Architecture for Time Series Forecasting.\"](http://arxiv.org/abs/2303.06053)<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Figure 1. TSMixer for multivariate time series forecasting.](imgs_models/tsmixer.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "from fastcore.test import test_eq\n",
    "from nbdev.showdoc import show_doc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from neuralforecast.losses.pytorch import MAE\n",
    "from neuralforecast.common._base_multivariate import BaseMultivariate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Auxiliary Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.1 Mixing layers\n",
    "A mixing layer consists of a sequential time- and feature Multi Layer Perceptron (`MLP`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "class TemporalMixing(nn.Module):\n",
    "    def __init__(self, n_series, input_size, dropout):\n",
    "        super().__init__()\n",
    "        self.temporal_norm = nn.BatchNorm1d(num_features=n_series * input_size, eps=0.001, momentum=0.01)\n",
    "        self.temporal_lin = nn.Linear(input_size, input_size)\n",
    "        self.temporal_drop = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, input):\n",
    "        # Get shapes\n",
    "        batch_size = input.shape[0]\n",
    "        input_size = input.shape[1]\n",
    "        n_series = input.shape[2]\n",
    "\n",
    "        # Temporal MLP\n",
    "        x = input.permute(0, 2, 1)                                      # [B, L, N] -> [B, N, L]\n",
    "        x = x.reshape(batch_size, -1)                                   # [B, N, L] -> [B, N * L]\n",
    "        x = self.temporal_norm(x)                                       # [B, N * L] -> [B, N * L]\n",
    "        x = x.reshape(batch_size, n_series, input_size)                 # [B, N * L] -> [B, N, L]\n",
    "        x = F.relu(self.temporal_lin(x))                                # [B, N, L] -> [B, N, L]\n",
    "        x = x.permute(0, 2, 1)                                          # [B, N, L] -> [B, L, N]\n",
    "        x = self.temporal_drop(x)                                       # [B, L, N] -> [B, L, N]\n",
    "\n",
    "        return x + input \n",
    "\n",
    "class FeatureMixing(nn.Module):\n",
    "    def __init__(self, n_series, input_size, dropout, ff_dim):\n",
    "        super().__init__()\n",
    "        self.feature_norm = nn.BatchNorm1d(num_features=n_series * input_size, eps=0.001, momentum=0.01)\n",
    "        self.feature_lin_1 = nn.Linear(n_series, ff_dim)\n",
    "        self.feature_lin_2 = nn.Linear(ff_dim, n_series)\n",
    "        self.feature_drop_1 = nn.Dropout(dropout)\n",
    "        self.feature_drop_2 = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, input):\n",
    "        # Get shapes\n",
    "        batch_size = input.shape[0]\n",
    "        input_size = input.shape[1]\n",
    "        n_series = input.shape[2]\n",
    "\n",
    "        # Feature MLP\n",
    "        x = input.reshape(batch_size, -1)                               # [B, L, N] -> [B, L * N]\n",
    "        x = self.feature_norm(x)                                        # [B, L * N] -> [B, L * N]\n",
    "        x = x.reshape(batch_size, input_size, n_series)                 # [B, L * N] -> [B, L, N]\n",
    "        x = F.relu(self.feature_lin_1(x))                               # [B, L, N] -> [B, L, ff_dim]\n",
    "        x = self.feature_drop_1(x)                                      # [B, L, ff_dim] -> [B, L, ff_dim]\n",
    "        x = self.feature_lin_2(x)                                       # [B, L, ff_dim] -> [B, L, N]\n",
    "        x = self.feature_drop_2(x)                                      # [B, L, N] -> [B, L, N]\n",
    "\n",
    "        return x + input \n",
    "\n",
    "class MixingLayer(nn.Module):\n",
    "    def __init__(self, n_series, input_size, dropout, ff_dim):\n",
    "        super().__init__()\n",
    "        # Mixing layer consists of a temporal and feature mixer\n",
    "        self.temporal_mixer = TemporalMixing(n_series, input_size, dropout)\n",
    "        self.feature_mixer = FeatureMixing(n_series, input_size, dropout, ff_dim)\n",
    "\n",
    "    def forward(self, input):\n",
    "        x = self.temporal_mixer(input)\n",
    "        x = self.feature_mixer(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.2 Reversible InstanceNormalization\n",
    "An Instance Normalization Layer that is reversible, based on [this reference implementation](https://github.com/google-research/google-research/blob/master/tsmixer/tsmixer_basic/models/rev_in.py).<br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "class ReversibleInstanceNorm1d(nn.Module):\n",
    "    def __init__(self, n_series, eps=1e-5):\n",
    "        super().__init__()\n",
    "        self.weight = nn.Parameter(torch.ones((1, 1, n_series)))\n",
    "        self.bias = nn.Parameter(torch.zeros((1, 1, n_series)))\n",
    "\n",
    "        self.eps = eps\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Batch statistics\n",
    "        self.batch_mean = torch.mean(x, axis=1, keepdim=True).detach()\n",
    "        self.batch_std = torch.sqrt(torch.var(x, axis=1, keepdim=True, unbiased=False) + self.eps).detach()\n",
    "        \n",
    "        # Instance normalization\n",
    "        x = x - self.batch_mean\n",
    "        x = x / self.batch_std\n",
    "        x = x * self.weight\n",
    "        x = x + self.bias\n",
    "        \n",
    "        return x\n",
    "\n",
    "    def reverse(self, x):\n",
    "        # Reverse the normalization\n",
    "        x = x - self.bias\n",
    "        x = x / self.weight       \n",
    "        x = x * self.batch_std\n",
    "        x = x + self.batch_mean       \n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class TSMixer(BaseMultivariate):\n",
    "    \"\"\" TSMixer\n",
    "\n",
    "    Time-Series Mixer (`TSMixer`) is a MLP-based multivariate time-series forecasting model. `TSMixer` jointly learns temporal and cross-sectional representations of the time-series by repeatedly combining time- and feature information using stacked mixing layers. A mixing layer consists of a sequential time- and feature Multi Layer Perceptron (`MLP`).\n",
    "\n",
    "    **Parameters:**<br>\n",
    "    `h`: int, forecast horizon.<br>\n",
    "    `input_size`: int, considered autorregresive inputs (lags), y=[1,2,3,4] input_size=2 -> lags=[1,2].<br>\n",
    "    `n_series`: int, number of time-series.<br>\n",
    "    `futr_exog_list`: str list, future exogenous columns.<br>\n",
    "    `hist_exog_list`: str list, historic exogenous columns.<br>\n",
    "    `stat_exog_list`: str list, static exogenous columns.<br>\n",
    "    `n_block`: int=2, number of mixing layers in the model.<br>\n",
    "    `ff_dim`: int=64, number of units for the second feed-forward layer in the feature MLP.<br>\n",
    "    `dropout`: float=0.9, dropout rate between (0, 1) .<br>\n",
    "    `revin`: bool=True, if True uses Reverse Instance Normalization to process inputs and outputs.<br>\n",
    "    `loss`: PyTorch module, instantiated train loss class from [losses collection](https://nixtla.github.io/neuralforecast/losses.pytorch.html).<br>\n",
    "    `valid_loss`: PyTorch module=`loss`, instantiated valid loss class from [losses collection](https://nixtla.github.io/neuralforecast/losses.pytorch.html).<br>\n",
    "    `max_steps`: int=1000, maximum number of training steps.<br>\n",
    "    `learning_rate`: float=1e-3, Learning rate between (0, 1).<br>\n",
    "    `num_lr_decays`: int=-1, Number of learning rate decays, evenly distributed across max_steps.<br>\n",
    "    `early_stop_patience_steps`: int=-1, Number of validation iterations before early stopping.<br>\n",
    "    `val_check_steps`: int=100, Number of training steps between every validation loss check.<br>\n",
    "    `batch_size`: int=32, number of different series in each batch.<br>\n",
    "    `step_size`: int=1, step size between each window of temporal data.<br>\n",
    "    `scaler_type`: str='identity', type of scaler for temporal inputs normalization see [temporal scalers](https://nixtla.github.io/neuralforecast/common.scalers.html).<br>\n",
    "    `random_seed`: int=1, random_seed for pytorch initializer and numpy generators.<br>\n",
    "    `num_workers_loader`: int=os.cpu_count(), workers to be used by `TimeSeriesDataLoader`.<br>\n",
    "    `drop_last_loader`: bool=False, if True `TimeSeriesDataLoader` drops last non-full batch.<br>\n",
    "    `alias`: str, optional,  Custom name of the model.<br>\n",
    "    `optimizer`: Subclass of 'torch.optim.Optimizer', optional, user specified optimizer instead of the default choice (Adam).<br>\n",
    "    `optimizer_kwargs`: dict, optional, list of parameters used by the user specified `optimizer`.<br>\n",
    "    `lr_scheduler`: Subclass of 'torch.optim.lr_scheduler.LRScheduler', optional, user specified lr_scheduler instead of the default choice (StepLR).<br>\n",
    "    `lr_scheduler_kwargs`: dict, optional, list of parameters used by the user specified `lr_scheduler`.<br>    \n",
    "    `**trainer_kwargs`: int,  keyword trainer arguments inherited from [PyTorch Lighning's trainer](https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.trainer.trainer.Trainer.html?highlight=trainer).<br>    \n",
    "\n",
    "    **References:**<br>\n",
    "    - [Chen, Si-An, Chun-Liang Li, Nate Yoder, Sercan O. Arik, and Tomas Pfister (2023). \"TSMixer: An All-MLP Architecture for Time Series Forecasting.\"](http://arxiv.org/abs/2303.06053)\n",
    "\n",
    "    \"\"\"\n",
    "    # Class attributes\n",
    "    SAMPLING_TYPE = 'multivariate'\n",
    "    EXOGENOUS_FUTR = False\n",
    "    EXOGENOUS_HIST = False\n",
    "    EXOGENOUS_STAT = False\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 input_size,\n",
    "                 n_series,\n",
    "                 futr_exog_list = None,\n",
    "                 hist_exog_list = None,\n",
    "                 stat_exog_list = None,\n",
    "                 n_block = 2,\n",
    "                 ff_dim = 64,\n",
    "                 dropout = 0.9,\n",
    "                 revin = True,\n",
    "                 loss = MAE(),\n",
    "                 valid_loss = None,\n",
    "                 max_steps: int = 1000,\n",
    "                 learning_rate: float = 1e-3,\n",
    "                 num_lr_decays: int = -1,\n",
    "                 early_stop_patience_steps: int =-1,\n",
    "                 val_check_steps: int = 100,\n",
    "                 batch_size: int = 32,\n",
    "                 step_size: int = 1,\n",
    "                 scaler_type: str = 'identity',\n",
    "                 random_seed: int = 1,\n",
    "                 num_workers_loader: int = 0,\n",
    "                 drop_last_loader: bool = False,\n",
    "                 optimizer = None,\n",
    "                 optimizer_kwargs = None,\n",
    "                 lr_scheduler = None,\n",
    "                 lr_scheduler_kwargs = None,\n",
    "                 **trainer_kwargs):\n",
    "\n",
    "        # Inherit BaseMultivariate class\n",
    "        super(TSMixer, self).__init__(h=h,\n",
    "                                    input_size=input_size,\n",
    "                                    n_series=n_series,\n",
    "                                    futr_exog_list=futr_exog_list,\n",
    "                                    hist_exog_list=hist_exog_list,\n",
    "                                    stat_exog_list=stat_exog_list,\n",
    "                                    loss=loss,\n",
    "                                    valid_loss=valid_loss,\n",
    "                                    max_steps=max_steps,\n",
    "                                    learning_rate=learning_rate,\n",
    "                                    num_lr_decays=num_lr_decays,\n",
    "                                    early_stop_patience_steps=early_stop_patience_steps,\n",
    "                                    val_check_steps=val_check_steps,\n",
    "                                    batch_size=batch_size,\n",
    "                                    step_size=step_size,\n",
    "                                    scaler_type=scaler_type,\n",
    "                                    random_seed=random_seed,\n",
    "                                    num_workers_loader=num_workers_loader,\n",
    "                                    drop_last_loader=drop_last_loader,\n",
    "                                    optimizer=optimizer,\n",
    "                                    optimizer_kwargs=optimizer_kwargs,\n",
    "                                    lr_scheduler=lr_scheduler,\n",
    "                                    lr_scheduler_kwargs=lr_scheduler_kwargs,\n",
    "                                    **trainer_kwargs)\n",
    "\n",
    "        # Reversible InstanceNormalization layer\n",
    "        self.revin = revin\n",
    "        if self.revin:\n",
    "            self.norm = ReversibleInstanceNorm1d(n_series = n_series)\n",
    "\n",
    "        # Mixing layers\n",
    "        mixing_layers = [MixingLayer(n_series=n_series, \n",
    "                                     input_size=input_size, \n",
    "                                     dropout=dropout, \n",
    "                                     ff_dim=ff_dim) \n",
    "                                     for _ in range(n_block)]\n",
    "        self.mixing_layers = nn.Sequential(*mixing_layers)\n",
    "\n",
    "        # Linear output with Loss dependent dimensions\n",
    "        self.out = nn.Linear(in_features=input_size, \n",
    "                             out_features=h * self.loss.outputsize_multiplier)\n",
    "\n",
    "    def forward(self, windows_batch):\n",
    "        # Parse batch\n",
    "        x = windows_batch['insample_y']  # x: [batch_size, input_size, n_series]\n",
    "        batch_size = x.shape[0]\n",
    "\n",
    "        # TSMixer: InstanceNorm + Mixing layers + Dense output layer + ReverseInstanceNorm\n",
    "        if self.revin:\n",
    "            x = self.norm(x)\n",
    "        x = self.mixing_layers(x)\n",
    "        x = x.permute(0, 2, 1)\n",
    "        x = self.out(x)\n",
    "        x = x.permute(0, 2, 1)\n",
    "        if self.revin:\n",
    "            x = self.norm.reverse(x)\n",
    "\n",
    "        x = x.reshape(batch_size, self.h, self.loss.outputsize_multiplier * self.n_series)\n",
    "        forecast = self.loss.domain_map(x)\n",
    "\n",
    "        # domain_map might have squeezed the last dimension in case n_series == 1\n",
    "        # Note that this fails in case of a tuple loss, but Multivariate does not support tuple losses yet.\n",
    "        if forecast.ndim == 2:\n",
    "            return forecast.unsqueeze(-1)\n",
    "        else:\n",
    "            return forecast"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(TSMixer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(TSMixer.fit, name='TSMixer.fit')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(TSMixer.predict, name='TSMixer.predict')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "import logging\n",
    "import warnings\n",
    "\n",
    "from neuralforecast import NeuralForecast\n",
    "from neuralforecast.utils import AirPassengersPanel, AirPassengersStatic\n",
    "from neuralforecast.losses.pytorch import MAE, MSE, RMSE, MAPE, SMAPE, MASE, relMSE, QuantileLoss, MQLoss, DistributionLoss,PMM, GMM, NBMM, HuberLoss, TukeyLoss, HuberQLoss, HuberMQLoss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Test losses\n",
    "logging.getLogger(\"pytorch_lightning\").setLevel(logging.ERROR)\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "Y_train_df = AirPassengersPanel[AirPassengersPanel.ds<AirPassengersPanel['ds'].values[-12]].reset_index(drop=True) # 132 train\n",
    "Y_test_df = AirPassengersPanel[AirPassengersPanel.ds>=AirPassengersPanel['ds'].values[-12]].reset_index(drop=True) # 12 test\n",
    "\n",
    "AirPassengersStatic_single = AirPassengersStatic[AirPassengersStatic[\"unique_id\"] == 'Airline1']\n",
    "Y_train_df_single = Y_train_df[Y_train_df[\"unique_id\"] == 'Airline1']\n",
    "Y_test_df_single = Y_test_df[Y_test_df[\"unique_id\"] == 'Airline1']\n",
    "\n",
    "losses = [MAE(), MSE(), RMSE(), MAPE(), SMAPE(), MASE(seasonality=12), relMSE(y_train=Y_train_df), QuantileLoss(q=0.5), MQLoss(), DistributionLoss(distribution='Bernoulli'), DistributionLoss(distribution='Normal'), DistributionLoss(distribution='Poisson'), DistributionLoss(distribution='StudentT'), DistributionLoss(distribution='NegativeBinomial'), DistributionLoss(distribution='Tweedie'), PMM(), GMM(), NBMM(), HuberLoss(), TukeyLoss(), HuberQLoss(q=0.5), HuberMQLoss()]\n",
    "valid_losses = [MAE(), MSE(), RMSE(), MAPE(), SMAPE(), MASE(seasonality=12), relMSE(y_train=Y_train_df), QuantileLoss(q=0.5), MQLoss(), DistributionLoss(distribution='Bernoulli'), DistributionLoss(distribution='Normal'), DistributionLoss(distribution='Poisson'), DistributionLoss(distribution='StudentT'), DistributionLoss(distribution='NegativeBinomial'), DistributionLoss(distribution='Tweedie'), PMM(), GMM(), NBMM(), HuberLoss(), TukeyLoss(), HuberQLoss(q=0.5), HuberMQLoss()]\n",
    "\n",
    "for loss, valid_loss in zip(losses, valid_losses):\n",
    "    try:\n",
    "        model = TSMixer(h=12,\n",
    "                        input_size=24,\n",
    "                        n_series=2,\n",
    "                        n_block=4,\n",
    "                        ff_dim=4,\n",
    "                        revin=True,\n",
    "                        scaler_type='standard',\n",
    "                        max_steps=2,\n",
    "                        early_stop_patience_steps=-1,\n",
    "                        val_check_steps=5,\n",
    "                        learning_rate=1e-3,\n",
    "                        loss=loss,\n",
    "                        valid_loss=valid_loss,\n",
    "                        batch_size=32\n",
    "                        )\n",
    "\n",
    "        fcst = NeuralForecast(models=[model], freq='M')\n",
    "        fcst.fit(df=Y_train_df, static_df=AirPassengersStatic, val_size=12)\n",
    "        forecasts = fcst.predict(futr_df=Y_test_df)\n",
    "    except Exception as e:\n",
    "        assert str(e) == f\"{loss} is not supported in a Multivariate model.\"\n",
    "\n",
    "\n",
    "# Test n_series = 1\n",
    "model = TSMixer(h=12,\n",
    "                input_size=24,\n",
    "                n_series=1,\n",
    "                n_block=4,\n",
    "                ff_dim=4,\n",
    "                revin=True,\n",
    "                scaler_type='standard',\n",
    "                max_steps=2,\n",
    "                early_stop_patience_steps=-1,\n",
    "                val_check_steps=5,\n",
    "                learning_rate=1e-3,\n",
    "                loss=MAE(),\n",
    "                valid_loss=MAE(),\n",
    "                batch_size=32\n",
    "                )\n",
    "fcst = NeuralForecast(models=[model], freq='M')\n",
    "fcst.fit(df=Y_train_df_single, static_df=AirPassengersStatic_single, val_size=12)\n",
    "forecasts = fcst.predict(futr_df=Y_test_df_single)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Usage Examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Train model and forecast future values with `predict` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| eval: false\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pytorch_lightning as pl\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from neuralforecast import NeuralForecast\n",
    "from neuralforecast.utils import AirPassengersPanel, AirPassengersStatic\n",
    "from neuralforecast.losses.pytorch import MAE\n",
    "\n",
    "Y_train_df = AirPassengersPanel[AirPassengersPanel.ds<AirPassengersPanel['ds'].values[-12]].reset_index(drop=True) # 132 train\n",
    "Y_test_df = AirPassengersPanel[AirPassengersPanel.ds>=AirPassengersPanel['ds'].values[-12]].reset_index(drop=True) # 12 test\n",
    "\n",
    "model = TSMixer(h=12,\n",
    "                input_size=24,\n",
    "                n_series=2, \n",
    "                n_block=4,\n",
    "                ff_dim=4,\n",
    "                dropout=0,\n",
    "                revin=True,\n",
    "                scaler_type='standard',\n",
    "                max_steps=200,\n",
    "                early_stop_patience_steps=-1,\n",
    "                val_check_steps=5,\n",
    "                learning_rate=1e-3,\n",
    "                loss=MAE(),\n",
    "                valid_loss=MAE(),\n",
    "                batch_size=32\n",
    "                )\n",
    "\n",
    "fcst = NeuralForecast(models=[model], freq='M')\n",
    "fcst.fit(df=Y_train_df, static_df=AirPassengersStatic, val_size=12)\n",
    "forecasts = fcst.predict(futr_df=Y_test_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| eval: false\n",
    "# Plot predictions\n",
    "fig, ax = plt.subplots(1, 1, figsize = (20, 7))\n",
    "Y_hat_df = forecasts.reset_index(drop=False).drop(columns=['unique_id','ds'])\n",
    "plot_df = pd.concat([Y_test_df, Y_hat_df], axis=1)\n",
    "plot_df = pd.concat([Y_train_df, plot_df])\n",
    "\n",
    "plot_df = plot_df[plot_df.unique_id=='Airline1'].drop('unique_id', axis=1)\n",
    "plt.plot(plot_df['ds'], plot_df['y'], c='black', label='True')\n",
    "plt.plot(plot_df['ds'], plot_df['TSMixer'], c='blue', label='Forecast')\n",
    "ax.set_title('AirPassengers Forecast', fontsize=22)\n",
    "ax.set_ylabel('Monthly Passengers', fontsize=20)\n",
    "ax.set_xlabel('Year', fontsize=20)\n",
    "ax.legend(prop={'size': 15})\n",
    "ax.grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using `cross_validation` to forecast multiple historic values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| eval: false\n",
    "fcst = NeuralForecast(models=[model], freq='M')\n",
    "forecasts = fcst.cross_validation(df=AirPassengersPanel, static_df=AirPassengersStatic, n_windows=2, step_size=12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| eval: false\n",
    "# Plot predictions\n",
    "fig, ax = plt.subplots(1, 1, figsize = (20, 7))\n",
    "Y_hat_df = forecasts.loc['Airline1']\n",
    "Y_df = AirPassengersPanel[AirPassengersPanel['unique_id']=='Airline1']\n",
    "\n",
    "plt.plot(Y_df['ds'], Y_df['y'], c='black', label='True')\n",
    "plt.plot(Y_hat_df['ds'], Y_hat_df['TSMixer'], c='blue', label='Forecast')\n",
    "ax.set_title('AirPassengers Forecast', fontsize=22)\n",
    "ax.set_ylabel('Monthly Passengers', fontsize=20)\n",
    "ax.set_xlabel('Year', fontsize=20)\n",
    "ax.legend(prop={'size': 15})\n",
    "ax.grid()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
